-module(mysql_protocol).
-export([init/1, send/2, recv/0, recv_result/0]).
-export([get_string/1, get_binary/1, get_length/1, write_length/1, write_binary/1, new_password/2]).
-include("mysql_define.hrl").

init(Socket) ->
	gen_tcp:controlling_process(Socket, self()),
	inet:setopts(Socket, [binary,{active, false}]),
	put(socket, Socket),
	put(recvbuf, <<>>).

send(Id, Packet) ->
	Sock = get(socket),
	%io:format("send packet, id:~p, length:~p~n", [Id, size(Packet)]),

	Len = size(Packet),
	Bin = <<Len:?UINT24, Id:?UINT8, Packet/binary>>,
	gen_tcp:send(Sock, Bin).

recv() ->
	Buf = get(recvbuf),
	Socket = get(socket),

	case Buf of
		<<Len:?UINT24, Id:?UINT8, Packet:Len/binary, Rest/binary>> -> 
			%io:format("recv packet, id:~p, length:~p~n", [Id, size(Packet)]),
			put(recvbuf, Rest),
			{ok, Id, Packet};
		_ ->
			case gen_tcp:recv(Socket, 0) of 
				{ok, Bin} ->
					put(recvbuf, iolist_to_binary([Buf, Bin])),
					recv();
				Err ->
					Err
			end
	end.

get_binary(B) ->
	{Len, B1} = get_length(B),
	{binary:part(B1, 0, Len), binary:part(B1, Len, size(B1) - Len)}.

get_string(B) ->
	get_string(B, 0).

get_string(B, Idx) ->
	case binary:at(B, Idx) of
		0 ->
			{binary_to_list(binary:part(B, 0, Idx)), binary:part(B, Idx + 1, size(B) - Idx - 1)};
		_ ->
			get_string(B, Idx + 1)
	end.


get_length(B) -> 
	<<Tag:?UINT8, Rest/binary>> = B,
	case Tag of
		Tag when Tag =< 250 ->
			{Tag, Rest};
		251 ->
			{null, Rest};
		252 ->
			<<Varint:?UINT16, Rest1/binary>> = Rest,
			{Varint, Rest1};
		253 ->
			<<Varint:?UINT32, Rest1/binary>> = Rest,
			{Varint, Rest1};
		254 -> 
			<<Varint:?UINT64, Rest1/binary>> = Rest,
			{Varint, Rest1};
		_ ->
			{badarg}
	end.

write_length(L) ->
	case L of
		L when L =< 250 ->
			<<L:?UINT8>>;
		L when L =< 65536 ->
			<<252, L:?UINT16>>;
		L when L =< 4294967296 ->
			<<253, L:?UINT32>>;
		_ ->
			<<254, L:?UINT64>>
	end.

write_binary(null) -> <<251>>;
write_binary(B) ->
	L = write_length( size(B) ),
	<<L/binary, B/binary>>.

dualmap(_F, [], []) -> [];
dualmap(F, [E1 | R1], [E2 | R2]) -> [F(E1, E2) | dualmap(F, R1, R2)].

new_password(Password, Salt) ->
	Stage1 = crypto:sha(Password),
	Stage2 = crypto:sha(Stage1),
	Res = crypto:sha_final(
			crypto:sha_update(
				crypto:sha_update(crypto:sha_init(), Salt),
				Stage2)
			),
	bxor_binary(Res, Stage1).

bxor_binary(B1, B2) ->
	list_to_binary(dualmap(fun (E1, E2) -> E1 bxor E2 end, binary_to_list(B1), binary_to_list(B2))).

recv_result() ->
	{ok, _, Packet} = recv(),
	<<FieldCount:?UINT8, B0/binary>> = Packet,
	%io:format("recv result packet: field_count:~p length:~p~n", [FieldCount, size(Packet)]),

	case FieldCount of
		255 -> % error
			<<ErrorNumber:?UINT16, SqlState:6/binary, ErrorMessage/binary>> = B0,
			%io:format("errcode:~p, state:~p, errmsg:~p~n", [ErrorNumber, SqlState, ErrorMessage]),
			{error, ErrorNumber, SqlState, binary_to_list(ErrorMessage)};
        254 -> % end of result set
			<<_WarningCount:?UINT16, _StatusFlags:?UINT16>> = B0,
			{eof, _WarningCount, _StatusFlags};
        0 -> % update success,
			{AffectedRows, B1} = get_length(B0),
			{InsertId, B2} = get_length(B1),
			<<_ServerStatus:?UINT16, _WarningCount:?UINT16, _Message/binary>> = B2,
			{ok, FieldCount, AffectedRows, InsertId};
		_ -> % query success
			{resultset, FieldCount}
	end.

